#include <iostream>
#include <armadillo>
#include <boost/program_options.hpp>

namespace po = boost::program_options;
extern po::variables_map vm;

#define ENDL std::endl
#define pb push_back
#define sz(c) ((int)(c).size())
#define nele(c) ((int)(c).n_elem)
#define tr(c, it) for (typeof((c).begin()) it = (c).begin(); it != (c).end(); it++)
typedef std::vector<int> vi_t;
typedef std::vector<vi_t> vvi_t;
typedef arma::vec vd_t;
typedef std::vector<vd_t> vvd_t;





struct Anneal {
  int n_base_, motif_length_;

  Anneal(int n_base, int motif_length) : n_base_(n_base), motif_length_(motif_length) {}

  void generate_freq(vvd_t& freq);
  void generate_data(int n_seq, int seq_length, const vvd_t& freq, vvi_t& data, vi_t& motif_pos);
  void generate_pos(int n_seq, int seq_length, vi_t& motif_pos);

  void print_freq(const vvd_t& freq);
  void print_data(const vvi_t& data, const vi_t& motif_pos);

  double score(const vvi_t& data, const vi_t& motif_pos);
  void compute_freq(const vvi_t& data, const vi_t& motif_pos, vvd_t& freq);
  double compute_score(const vvi_t& data, const vi_t& motif_pos, const vvd_t& freq);
  
  void anneal(const vvi_t& data, const vi_t& start_pos, vi_t& end_pos, vvd_t& end_freq);
  
  static int sample(const vd_t& f);
};

void Anneal::generate_freq(vvd_t& freq) 
{
  freq.clear();
  double f[][4] = {{0.25, 0.25, 0.25, 0.25}, 
		   {0.7, 0.1, 0.1, 0.1}, 
		   {0.7, 0.1, 0.1, 0.1},
		   {0.1, 0.7, 0.1, 0.1},
		   {0.1, 0.1, 0.7, 0.1},
		   {0.1, 0.1, 0.1, 0.7},
		   {0.1, 0.1, 0.1, 0.7},
  };
  for (int i = 0; i <= motif_length_; i++) {
    vd_t vf(n_base_);
    for (int j = 0; j < n_base_; j++) vf[j] = f[i][j];
    vf /= arma::sum(vf);
    freq.pb(vf);
  }
}

void Anneal::generate_data(int n_seq, int seq_length, const vvd_t& freq, vvi_t& data, vi_t& motif_pos)
{
  data.clear();
  motif_pos.clear();
  int motif_length = sz(freq) - 1;
  for (int i = 0; i < n_seq; i++) {
    int max_pos = seq_length - motif_length;
    // random motif position
    int pos = rand() % max_pos;
    motif_pos.pb(pos);
    // random bases on non-motif and motif position
    vi_t d(seq_length);
    for (int p = 0; p < seq_length; p++) {
      if (p < pos || p >= pos+motif_length) 
	d[p] = sample(freq[0]);
      else
	d[p] = sample(freq[p-pos+1]);
    }
    data.pb(d);
  }
}

void Anneal::generate_pos(int n_seq, int seq_length, vi_t& motif_pos)
{
  motif_pos.clear();
  for (int i = 0; i < n_seq; i++) {
    int max_pos = seq_length - motif_length_;
    // random motif position
    int pos = rand() % max_pos;
    motif_pos.pb(pos);
  }
}

int Anneal::sample(const vd_t& f) 
{
  double s = arma::sum(f), ss = 0;
  double r = arma::randu()*s;
  for (int i = 0; i < (int) f.n_elem; i++) {
    ss += f[i];
    if (ss >= r) return i;
  }
  return -1;
}

void Anneal::print_freq(const vvd_t& freq) 
{
  for (int i = 0; i < sz(freq); i++) {
    std::cout << i << ":";
    for (int j = 0; j < nele(freq[i]); j++) {
      std::cout << " " << freq[i][j];
    }
    std::cout << ENDL;
  }
}

void Anneal::print_data(const vvi_t& data, const vi_t& motif_pos)
{
  int max_len = 0;
  tr(data, it) if (sz(*it) > max_len) max_len = sz(*it);
  int i = 0;
  tr(data, it) {
    int pos = motif_pos[i];
    const vi_t& d = data[i];

    for (int j = 0; j < max_len - pos; j++) std::cout << " ";
    for (int j = 0; j < pos; j++) std::cout << d[j];
    std::cout << "  ";
    for (int j = pos; j < pos + motif_length_; j++) std::cout << d[j];
    std::cout << "  ";
    for (int j = pos + motif_length_; j < sz(d); j++) std::cout << d[j];
    std::cout <<ENDL;

    ++i;
  }
}

double Anneal::score(const vvi_t& data, const vi_t& motif_pos)
{
  vvd_t freq;
  compute_freq(data, motif_pos, freq);

  double s = compute_score(data, motif_pos, freq);
  return s;
}

void Anneal::compute_freq(const vvi_t& data, const vi_t& motif_pos, vvd_t& freq)
{
  freq.clear();
  for (int i = 0; i <= motif_length_; i++) freq.pb(arma::zeros<vd_t>(n_base_));
  
  int n_seq = sz(data);
  for (int i = 0; i < n_seq; i++) { // count occurence of non-motif and motif position
    int seq_len = sz(data[i]);
    int pos = motif_pos[i];
    for (int p = 0; p < seq_len; p++) {
      int d = data[i][p];
      if (p < pos || p >= pos+motif_length_) 
	freq[0][d] += 1.0;
      else
	freq[p-pos+1][d] += 1.0;
    }
  }
  tr(freq, it) *it /= arma::sum(*it); // normalize
  //std::cout << "TEST\n"; print_freq(freq);
}

double Anneal::compute_score(const vvi_t& data, const vi_t& motif_pos, const vvd_t& freq)
{
  double score = 0;
  for (int i = 0; i < sz(data); i++) {
    int pos = motif_pos[i];
    double s = 0;
    for (int j = 0; j < motif_length_; j++) {
      int d = data[i][pos+j];
      s += log(freq[j+1][d] / freq[0][d]);
    }    
    score += s;
  } 
  return score / sz(data); // normalize 
}

void Anneal::anneal(const vvi_t& data, const vi_t& start_pos, vi_t& end_pos, vvd_t& end_freq)
{
  int T = 20000;
  int n_seq = sz(data);

  vi_t current_pos = start_pos;
  double current_score = score(data, current_pos);
  double max_temp = .001;

  for (int iter = 0; iter < T; iter++) {
    double temp = (double)(T-iter)/T*max_temp;  // temperature
    double prob = 0;
    double dscore = 0;
    int choose = 0;

    int pick_seq = rand() % n_seq;
    
    int seq_len = sz(data[pick_seq]);
    vi_t pos = current_pos, max_pos;
    double max_score = -INFINITY;
    
    for (int j = 0; j < seq_len - motif_length_; j++) {
      pos[pick_seq] = rand() % (seq_len - motif_length_);
      double pos_score = score(data, pos);
      if (max_score < pos_score) {
	max_score = pos_score;
	max_pos = pos;
      }
    }
    
    if (max_score > current_score) {   // always go up hill
      current_pos = max_pos;
      current_score = max_score;
    }
    else {                             // accept down hill step with some probability
      dscore = current_score - max_score; // dscore >= 0
      double r = arma::randu();                  // a random number
      prob = exp(-dscore/temp);

      if (prob > r) {
	current_pos = max_pos;
	current_score = max_score;
	choose = 1;
      }
    }
    std::cout << "(iter=" << iter << " score=" << max_score 
	      << " temp=" << temp << " dscore=" << dscore 
	      << " prob=" << prob << " " << choose << ") "; 
    std::cout.flush();
    /*if (iter % 4 == 3)*/ std::cout << ENDL;

    std::cerr << "(iter=" << iter << " score=" << max_score 
	      << " temp=" << temp << " dscore=" << dscore 
	      << " prob=" << prob << " " << choose << ") "; 
    std::cerr.flush();
    /*if (iter % 4 == 3)*/ std::cerr << ENDL;
  }

  end_pos = current_pos;
  compute_freq(data, end_pos, end_freq);
}

void anneal_test()
{
  std::cout << "================ ANNEAL ================ " << ENDL;
  int n_base = 4;
  int motif_length = 6;
  int n_seq = 1000;
  int seq_length = 30;

  vvd_t freq, anneal_freq;
  vvi_t data;
  vi_t motif_pos, rand_pos, anneal_pos;

  Anneal a(n_base, motif_length);

  a.generate_freq(freq);
  a.generate_data(n_seq, seq_length, freq, data, motif_pos);

  a.print_data(data, motif_pos);
  a.print_freq(freq);

  std::cout << "score(true_motif)=" << a.score(data, motif_pos) << ENDL;

  a.generate_pos(n_seq, seq_length, rand_pos);
  std::cout << "score(rand_motif)=" << a.score(data, rand_pos) << ENDL;

  a.anneal(data, rand_pos, anneal_pos, anneal_freq);
  a.print_data(data, anneal_pos);
  std::cout << "annealed frequency:\n"; a.print_freq(anneal_freq);

}
